package spark.product;

import org.apache.spark.sql.Row;
import org.apache.spark.sql.expressions.MutableAggregationBuffer;
import org.apache.spark.sql.expressions.UserDefinedAggregateFunction;
import org.apache.spark.sql.types.DataType;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.StructType;

import java.util.Arrays;

public class GroupConcatDistinctUDAF extends UserDefinedAggregateFunction {
    //输入数据的字段与类型
    private StructType inputSchema = DataTypes.createStructType(Arrays.asList(
            DataTypes.createStructField("city_id", DataTypes.StringType,true)
    ));
    //缓冲数据的字段与类型，中间聚合时处理的数据类型
    private StructType bufferSchema = DataTypes.createStructType(Arrays.asList(
            DataTypes.createStructField("buffercityInfo", DataTypes.StringType, true)
    ));

    @Override
    public StructType inputSchema() {
        return inputSchema;
    }

    @Override
    public StructType bufferSchema() {
        return bufferSchema;
    }

    @Override
    //指定返回类型
    public DataType dataType() {
        return DataTypes.StringType;
    }

    @Override
    //指定是否是确定性
    public boolean deterministic() {
        return true;
    }

    @Override
    //内部指定一个初始值
    public void initialize(MutableAggregationBuffer buffer) {
        buffer.update(0, "");
    }

    @Override
    public void update(MutableAggregationBuffer buffer, Row input) {
        // 缓冲中的已经拼接过的城市信息串
        String bufferCityInfo = buffer.getString(0);
        // 刚刚传递进来的某个城市信息
        String cityInfo = input.getString(0);

        // 在这里要实现去重的逻辑
        // 判断：之前没有拼接过某个城市信息，那么这里才可以接下去拼接新的城市信息
        if(!bufferCityInfo.contains(cityInfo)) {
            if("".equals(bufferCityInfo)) {
                bufferCityInfo += cityInfo;
            } else {
                // 比如1:北京
                // 1:北京,2:上海
                bufferCityInfo += "," + cityInfo;
            }

            buffer.update(0, bufferCityInfo);
        }
    }

    @Override
    public void merge(MutableAggregationBuffer buffer1, Row buffer2) {
        String bufferCityInfo1 = buffer1.getString(0);
        String bufferCityInfo2 = buffer2.getString(0);

        for(String cityInfo : bufferCityInfo2.split(",")) {
            if(!bufferCityInfo1.contains(cityInfo)) {
                if("".equals(bufferCityInfo1)) {
                    bufferCityInfo1 += cityInfo;
                } else {
                    bufferCityInfo1 += "," + cityInfo;
                }
            }
        }

        buffer1.update(0, bufferCityInfo1);

    }

    @Override
    public Object evaluate(Row buffer) {
        return buffer.getString(0);
    }
}
